{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "75619c50",
   "metadata": {},
   "source": [
    "# Sidekick with Planner and Clarification"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "612ee879",
   "metadata": {},
   "source": [
    "This enhanced version includes:\n",
    "- 🤔 Clarifying questions before starting work\n",
    "- 📋 Planner agent that creates execution plans\n",
    "- 🛠️ Full toolkit (browser, files, search, notifications, Python, Wikipedia)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3932e419",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Annotated, TypedDict, List, Dict, Any, Optional\n",
    "from langchain_core.messages import AIMessage, HumanMessage, SystemMessage\n",
    "from langchain_openai import ChatOpenAI\n",
    "from langgraph.graph import StateGraph, START, END\n",
    "from langgraph.checkpoint.memory import MemorySaver\n",
    "from langgraph.prebuilt import ToolNode\n",
    "from langgraph.graph.message import add_messages\n",
    "from pydantic import BaseModel, Field\n",
    "from IPython.display import Image, display\n",
    "import gradio as gr\n",
    "from dotenv import load_dotenv\n",
    "from datetime import datetime\n",
    "from playwright.async_api import async_playwright\n",
    "from langchain_community.agent_toolkits import PlayWrightBrowserToolkit\n",
    "import os\n",
    "import requests\n",
    "from langchain.agents import Tool\n",
    "from langchain_community.agent_toolkits import FileManagementToolkit\n",
    "from langchain_community.tools.wikipedia.tool import WikipediaQueryRun\n",
    "from langchain_experimental.tools import PythonREPLTool\n",
    "from langchain_community.utilities import GoogleSerperAPIWrapper\n",
    "from langchain_community.utilities.wikipedia import WikipediaAPIWrapper"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b783d47",
   "metadata": {},
   "outputs": [],
   "source": [
    "load_dotenv(override=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c565cb6",
   "metadata": {},
   "source": [
    "### Tools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b5c0a46c",
   "metadata": {},
   "outputs": [],
   "source": [
    "pushover_token = os.getenv(\"PUSHOVER_TOKEN\")\n",
    "pushover_user = os.getenv(\"PUSHOVER_USER\")\n",
    "pushover_url = \"https://api.pushover.net/1/messages.json\"\n",
    "serper = GoogleSerperAPIWrapper()\n",
    "\n",
    "async def playwright_tools():\n",
    "    playwright = await async_playwright().start()\n",
    "    browser = await playwright.chromium.launch(headless=False)\n",
    "    toolkit = PlayWrightBrowserToolkit.from_browser(async_browser=browser)\n",
    "    return toolkit.get_tools(), browser, playwright\n",
    "\n",
    "def push(text: str):\n",
    "    \"\"\"Send a push notification to the user\"\"\"\n",
    "    requests.post(pushover_url, data={\"token\": pushover_token, \"user\": pushover_user, \"message\": text})\n",
    "    return \"success\"\n",
    "\n",
    "def get_file_tools():\n",
    "    toolkit = FileManagementToolkit(root_dir=\"sandbox\")\n",
    "    return toolkit.get_tools()\n",
    "\n",
    "async def other_tools():\n",
    "    push_tool = Tool(name=\"send_push_notification\", func=push, description=\"Use this tool when you want to send a push notification\")\n",
    "    file_tools = get_file_tools()\n",
    "    tool_search = Tool(\n",
    "        name=\"search\",\n",
    "        func=serper.run,\n",
    "        description=\"Use this tool when you want to get the results of an online web search\"\n",
    "    )\n",
    "    wikipedia = WikipediaAPIWrapper()\n",
    "    wiki_tool = WikipediaQueryRun(api_wrapper=wikipedia)\n",
    "    python_repl = PythonREPLTool()\n",
    "    \n",
    "    return file_tools + [push_tool, tool_search, python_repl, wiki_tool]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "87a55cef",
   "metadata": {},
   "source": [
    "### Structured Outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6320e751",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ClarifyingQuestion(BaseModel):\n",
    "    question: str = Field(description=\"The clarifying question to ask\")\n",
    "    why_important: str = Field(description=\"Why this question matters\")\n",
    "\n",
    "class ClarifyingQuestions(BaseModel):\n",
    "    questions: List[ClarifyingQuestion] = Field(description=\"List of 2-4 clarifying questions\")\n",
    "\n",
    "class PlanStep(BaseModel):\n",
    "    step_number: int = Field(description=\"The step number\")\n",
    "    action: str = Field(description=\"What needs to be done in this step\")\n",
    "    tools_needed: List[str] = Field(description=\"List of tools that might be needed\")\n",
    "\n",
    "class ExecutionPlan(BaseModel):\n",
    "    steps: List[PlanStep] = Field(description=\"Ordered list of steps to complete the task\")\n",
    "    estimated_complexity: str = Field(description=\"Low, Medium, or High complexity\")\n",
    "\n",
    "class EvaluatorOutput(BaseModel):\n",
    "    feedback: str = Field(description=\"Feedback on the assistant's response\")\n",
    "    success_criteria_met: bool = Field(description=\"Whether the success criteria have been met\")\n",
    "    user_input_needed: bool = Field(description=\"True if more input is needed from the user\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9b5b1a14",
   "metadata": {},
   "source": [
    "### State"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ba96b15",
   "metadata": {},
   "outputs": [],
   "source": [
    "class State(TypedDict):\n",
    "    messages: Annotated[List[Any], add_messages]\n",
    "    success_criteria: str\n",
    "    clarifications: Optional[str]\n",
    "    execution_plan: Optional[str]\n",
    "    feedback_on_work: Optional[str]\n",
    "    success_criteria_met: bool\n",
    "    user_input_needed: bool\n",
    "    stage: str  # \"clarifying\", \"planning\", \"executing\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2d0754bf",
   "metadata": {},
   "source": [
    "### Setup and Configuration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e1f0bc0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Store everything in one dict\n",
    "state = {}\n",
    "\n",
    "async def setup():\n",
    "    \"\"\"Initialize tools and LLMs\"\"\"\n",
    "    # Setup tools\n",
    "    playwright_tools_list, state['browser'], state['playwright'] = await playwright_tools()\n",
    "    other_tools_list = await other_tools()\n",
    "    state['tools'] = playwright_tools_list + other_tools_list\n",
    "    \n",
    "    # Setup LLMs with structured outputs\n",
    "    worker_llm = ChatOpenAI(model=\"gpt-4o-mini\")\n",
    "    state['worker_llm_with_tools'] = worker_llm.bind_tools(state['tools'])\n",
    "    \n",
    "    clarifying_llm = ChatOpenAI(model=\"gpt-4o-mini\")\n",
    "    state['clarifying_llm'] = clarifying_llm.with_structured_output(ClarifyingQuestions)\n",
    "    \n",
    "    planner_llm = ChatOpenAI(model=\"gpt-4o-mini\")\n",
    "    state['planner_llm'] = planner_llm.with_structured_output(ExecutionPlan)\n",
    "    \n",
    "    evaluator_llm = ChatOpenAI(model=\"gpt-4o-mini\")\n",
    "    state['evaluator_llm'] = evaluator_llm.with_structured_output(EvaluatorOutput)\n",
    "    \n",
    "    state['memory'] = MemorySaver()\n",
    "    \n",
    "    await build_graph()\n",
    "\n",
    "def cleanup():\n",
    "    \"\"\"Cleanup browser resources\"\"\"\n",
    "    if 'browser' in state and state['browser']:\n",
    "        import asyncio\n",
    "        try:\n",
    "            loop = asyncio.get_running_loop()\n",
    "            loop.create_task(state['browser'].close())\n",
    "            if 'playwright' in state and state['playwright']:\n",
    "                loop.create_task(state['playwright'].stop())\n",
    "        except:\n",
    "            pass"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f86c3cc5",
   "metadata": {},
   "source": [
    "### Core Nodes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b98c986f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def clarifying_agent(graph_state: State) -> Dict[str, Any]:\n",
    "    \"\"\"Ask clarifying questions about the task\"\"\"\n",
    "    user_request = graph_state[\"messages\"][0].content\n",
    "    \n",
    "    prompt = f\"\"\"User's request: {user_request}\n",
    "\n",
    "    Success criteria: {graph_state['success_criteria']}\n",
    "\n",
    "    Generate 2-4 clarifying questions to better understand the user's needs.\n",
    "    Focus on scope, preferences, constraints, and desired output format.\"\"\"\n",
    "    \n",
    "    questions = state['clarifying_llm'].invoke([HumanMessage(content=prompt)])\n",
    "    \n",
    "    # Format questions\n",
    "    response = \"Before I start, I have some clarifying questions:\\n\\n\"\n",
    "    for i, q in enumerate(questions.questions, 1):\n",
    "        response += f\"**{i}. {q.question}**\\n\"\n",
    "    \n",
    "    return {\n",
    "        \"messages\": [AIMessage(content=response)],\n",
    "        \"stage\": \"planning\"\n",
    "    }\n",
    "\n",
    "def planner_agent(graph_state: State) -> Dict[str, Any]:\n",
    "    \"\"\"Create an execution plan for the task\"\"\"\n",
    "    user_request = graph_state[\"messages\"][0].content\n",
    "    clarifications = graph_state.get(\"clarifications\", \"No additional clarifications\")\n",
    "    \n",
    "    available_tools = \", \".join([tool.name for tool in state['tools']])\n",
    "    \n",
    "    prompt = f\"\"\"User's request: {user_request}\n",
    "\n",
    "    Clarifications: {clarifications}\n",
    "\n",
    "    Success criteria: {graph_state['success_criteria']}\n",
    "\n",
    "    Available tools: {available_tools}\n",
    "\n",
    "    Create a detailed execution plan with ordered steps to accomplish this task.\n",
    "    Consider which tools are needed for each step.\"\"\"\n",
    "        \n",
    "    plan = state['planner_llm'].invoke([HumanMessage(content=prompt)])\n",
    "    \n",
    "    # Format plan\n",
    "    plan_text = f\"Execution Plan (Complexity: {plan.estimated_complexity}):\\n\\n\"\n",
    "    for step in plan.steps:\n",
    "        plan_text += f\"Step {step.step_number}: {step.action}\\n\"\n",
    "        if step.tools_needed:\n",
    "            plan_text += f\"  Tools: {', '.join(step.tools_needed)}\\n\"\n",
    "    \n",
    "    return {\n",
    "        \"messages\": [AIMessage(content=plan_text)],\n",
    "        \"execution_plan\": plan_text,\n",
    "        \"stage\": \"executing\"\n",
    "    }\n",
    "\n",
    "def worker(graph_state: State) -> Dict[str, Any]:\n",
    "    \"\"\"Execute the plan using tools\"\"\"\n",
    "    system_message = f\"\"\"You are a helpful assistant executing a planned task.\n",
    "\n",
    "    Current date and time: {datetime.now().strftime(\"%Y-%m-%d %H:%M:%S\")}\n",
    "\n",
    "    Success criteria: {graph_state['success_criteria']}\n",
    "\n",
    "    Execution plan:\n",
    "    {graph_state.get('execution_plan', 'No formal plan')}\n",
    "\n",
    "    Work through the plan systematically. Use tools as needed.\n",
    "    When finished, provide a clear final response.\"\"\"\n",
    "    \n",
    "    if graph_state.get(\"feedback_on_work\"):\n",
    "        system_message += f\"\\n\\nPrevious feedback: {graph_state['feedback_on_work']}\\nPlease address this feedback.\"\n",
    "    \n",
    "    messages = [SystemMessage(content=system_message)] + graph_state[\"messages\"][1:]\n",
    "    response = state['worker_llm_with_tools'].invoke(messages)\n",
    "    \n",
    "    return {\"messages\": [response]}\n",
    "\n",
    "def evaluator(graph_state: State) -> Dict[str, Any]:\n",
    "    \"\"\"Evaluate if task is complete\"\"\"\n",
    "    last_response = graph_state[\"messages\"][-1].content\n",
    "    \n",
    "    prompt = f\"\"\"Evaluate this assistant's work:\n",
    "\n",
    "    Success criteria: {graph_state['success_criteria']}\n",
    "\n",
    "    Assistant's response: {last_response}\n",
    "\n",
    "    Decide if:\n",
    "    1. Success criteria is met\n",
    "    2. More user input is needed\n",
    "    3. Provide constructive feedback\"\"\"\n",
    "    \n",
    "    eval_result = state['evaluator_llm'].invoke([HumanMessage(content=prompt)])\n",
    "    \n",
    "    return {\n",
    "        \"messages\": [AIMessage(content=f\"Evaluation: {eval_result.feedback}\")],\n",
    "        \"feedback_on_work\": eval_result.feedback,\n",
    "        \"success_criteria_met\": eval_result.success_criteria_met,\n",
    "        \"user_input_needed\": eval_result.user_input_needed\n",
    "    }"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b84f5385",
   "metadata": {},
   "source": [
    "## Graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "697e85f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def worker_router(graph_state: State) -> str:\n",
    "    \"\"\"Route from worker to tools or evaluator\"\"\"\n",
    "    last_message = graph_state[\"messages\"][-1]\n",
    "    if hasattr(last_message, \"tool_calls\") and last_message.tool_calls:\n",
    "        return \"tools\"\n",
    "    return \"evaluator\"\n",
    "\n",
    "def route_after_evaluation(graph_state: State) -> str:\n",
    "    \"\"\"Route after evaluation - continue or end\"\"\"\n",
    "    if graph_state[\"success_criteria_met\"] or graph_state[\"user_input_needed\"]:\n",
    "        return \"END\"\n",
    "    return \"worker\"\n",
    "\n",
    "def route_from_start(graph_state: State) -> str:\n",
    "    \"\"\"Route based on current stage\"\"\"\n",
    "    stage = graph_state.get(\"stage\", \"clarifying\")\n",
    "    if stage == \"clarifying\":\n",
    "        return \"clarifying\"\n",
    "    elif stage == \"planning\":\n",
    "        return \"planning\"\n",
    "    return \"worker\"\n",
    "\n",
    "async def build_graph():\n",
    "    \"\"\"Build the LangGraph workflow\"\"\"\n",
    "    graph_builder = StateGraph(State)\n",
    "    \n",
    "    # Add nodes\n",
    "    graph_builder.add_node(\"clarifying\", clarifying_agent)\n",
    "    graph_builder.add_node(\"planning\", planner_agent)\n",
    "    graph_builder.add_node(\"worker\", worker)\n",
    "    graph_builder.add_node(\"tools\", ToolNode(tools=state['tools']))\n",
    "    graph_builder.add_node(\"evaluator\", evaluator)\n",
    "    \n",
    "    # Add edges\n",
    "    graph_builder.add_conditional_edges(\n",
    "        START, \n",
    "        route_from_start, \n",
    "        {\"clarifying\": \"clarifying\", \"planning\": \"planning\", \"worker\": \"worker\"}\n",
    "    )\n",
    "    graph_builder.add_edge(\"clarifying\", END)\n",
    "    graph_builder.add_edge(\"planning\", \"worker\")\n",
    "    graph_builder.add_conditional_edges(\n",
    "        \"worker\", \n",
    "        worker_router, \n",
    "        {\"tools\": \"tools\", \"evaluator\": \"evaluator\"}\n",
    "    )\n",
    "    graph_builder.add_edge(\"tools\", \"worker\")\n",
    "    graph_builder.add_conditional_edges(\n",
    "        \"evaluator\", \n",
    "        route_after_evaluation,\n",
    "        {\"worker\": \"worker\", \"END\": END}\n",
    "    )\n",
    "    \n",
    "    state['graph'] = graph_builder.compile(checkpointer=state['memory'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff6ed89c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import Image, display\n",
    "\n",
    "# First run setup to build the graph\n",
    "await setup()\n",
    "\n",
    "# Then display the graph visualization\n",
    "display(Image(state['graph'].get_graph().draw_mermaid_png()))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aae0fb12",
   "metadata": {},
   "source": [
    "### Chat Handler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d967715f",
   "metadata": {},
   "outputs": [],
   "source": [
    "async def chat(message: str, history: List, success_criteria: str):\n",
    "    \"\"\"Simplified chat handler\"\"\"\n",
    "    \n",
    "    if 'graph' not in state:\n",
    "        return history + [{\"role\": \"assistant\", \"content\": \"System initializing... Please run await setup() first.\"}]\n",
    "    \n",
    "    # Use a default thread ID\n",
    "    config = {\"configurable\": {\"thread_id\": \"default\"}}\n",
    "    \n",
    "    # Add user message\n",
    "    history = history + [{\"role\": \"user\", \"content\": message}]\n",
    "    \n",
    "    # Stage 1: Initial request -> Ask clarifying questions\n",
    "    if len(history) == 1:\n",
    "        graph_state = {\n",
    "            \"messages\": [HumanMessage(content=message)],\n",
    "            \"success_criteria\": success_criteria or \"Task completed accurately and thoroughly\",\n",
    "            \"stage\": \"clarifying\"\n",
    "        }\n",
    "        result = await state['graph'].ainvoke(graph_state, config=config)\n",
    "        return history + [{\"role\": \"assistant\", \"content\": result[\"messages\"][-1].content}]\n",
    "    \n",
    "    # Stage 2: Answer clarifying questions -> Create plan\n",
    "    elif \"clarifying questions\" in history[-2][\"content\"].lower():\n",
    "        original_request = history[0][\"content\"]\n",
    "        graph_state = {\n",
    "            \"messages\": [HumanMessage(content=original_request)],\n",
    "            \"success_criteria\": success_criteria or \"Task completed accurately and thoroughly\",\n",
    "            \"clarifications\": message,\n",
    "            \"stage\": \"planning\"\n",
    "        }\n",
    "        result = await state['graph'].ainvoke(graph_state, config=config)\n",
    "        return history + [{\"role\": \"assistant\", \"content\": result[\"messages\"][-1].content + \"\\n\\nType 'proceed' to start execution.\"}]\n",
    "    \n",
    "    # Stage 3: Review plan -> Execute\n",
    "    elif \"execution plan\" in history[-2][\"content\"].lower():\n",
    "        if \"proceed\" in message.lower():\n",
    "            original_request = history[0][\"content\"]\n",
    "            clarifications = history[2][\"content\"]\n",
    "            plan = history[-2][\"content\"]\n",
    "            \n",
    "            graph_state = {\n",
    "                \"messages\": [HumanMessage(content=original_request)],\n",
    "                \"success_criteria\": success_criteria or \"Task completed accurately and thoroughly\",\n",
    "                \"clarifications\": clarifications,\n",
    "                \"execution_plan\": plan,\n",
    "                \"stage\": \"executing\",\n",
    "                \"success_criteria_met\": False,\n",
    "                \"user_input_needed\": False\n",
    "            }\n",
    "            \n",
    "            result = await state['graph'].ainvoke(graph_state, config=config)\n",
    "            final_content = f\"{result['messages'][-2].content}\\n\\n---\\n{result['messages'][-1].content}\"\n",
    "            return history + [{\"role\": \"assistant\", \"content\": final_content}]\n",
    "        else:\n",
    "            return history + [{\"role\": \"assistant\", \"content\": \"Type 'proceed' when ready to execute.\"}]\n",
    "    \n",
    "    # Stage 4: New request\n",
    "    else:\n",
    "        graph_state = {\n",
    "            \"messages\": [HumanMessage(content=message)],\n",
    "            \"success_criteria\": success_criteria or \"Task completed accurately and thoroughly\",\n",
    "            \"stage\": \"clarifying\"\n",
    "        }\n",
    "        result = await state['graph'].ainvoke(graph_state, config=config)\n",
    "        return history + [{\"role\": \"assistant\", \"content\": result[\"messages\"][-1].content}]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "551cf92e",
   "metadata": {},
   "source": [
    "### UI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6316a09d",
   "metadata": {},
   "outputs": [],
   "source": [
    "with gr.Blocks(theme=gr.themes.Soft(primary_hue=\"blue\")) as demo:\n",
    "    gr.Markdown(\"# 🤖 Enhanced Sidekick\")\n",
    "    gr.Markdown(\"AI assistant with clarifying questions and execution planning\")\n",
    "    \n",
    "    chatbot = gr.Chatbot(label=\"Sidekick\", height=500, type=\"messages\")\n",
    "    \n",
    "    with gr.Row():\n",
    "        msg_input = gr.Textbox(\n",
    "            label=\"Message\",\n",
    "            placeholder=\"What would you like me to help with?\",\n",
    "            lines=2,\n",
    "            scale=4\n",
    "        )\n",
    "        send_btn = gr.Button(\"Send\", variant=\"primary\", scale=1)\n",
    "    \n",
    "    success_input = gr.Textbox(\n",
    "        label=\"Success Criteria (optional)\",\n",
    "        placeholder=\"What defines success for this task?\",\n",
    "        lines=1\n",
    "    )\n",
    "    \n",
    "    clear_btn = gr.Button(\"Clear Chat\", variant=\"secondary\")\n",
    "    \n",
    "    # Wire up events\n",
    "    send_btn.click(\n",
    "        chat,\n",
    "        inputs=[msg_input, chatbot, success_input],\n",
    "        outputs=[chatbot]\n",
    "    )\n",
    "    \n",
    "    msg_input.submit(\n",
    "        chat,\n",
    "        inputs=[msg_input, chatbot, success_input],\n",
    "        outputs=[chatbot]\n",
    "    )\n",
    "    \n",
    "    clear_btn.click(lambda: [], outputs=[chatbot])\n",
    "\n",
    "# Launch\n",
    "demo.launch(inbrowser=True)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
